import os
import shutil
import random
import xml.etree.ElementTree as ET


EXCLUDED_CLASSES = {"red"}


def count_classes_in_annotations(annotations_path, excluded_classes=None):
    if excluded_classes is None:
        excluded_classes = set()
    class_counts = {}
    for filename in os.listdir(annotations_path):
        if filename.endswith(".xml"):
            filepath = os.path.join(annotations_path, filename)
            tree = ET.parse(filepath)
            root = tree.getroot()
            for obj in root.findall("object"):
                name = obj.find("name").text
                if name in excluded_classes:
                    continue
                if name in class_counts:
                    class_counts[name] += 1
                else:
                    class_counts[name] = 1
    return class_counts


def merge_voc_folders(
    source_folders,
    target_folder,
    resampling_config,
    total_class_counts,
    excluded_classes=None,
):
    if excluded_classes is None:
        excluded_classes = set()
    # 创建目标文件夹及其子文件夹
    if not os.path.exists(os.path.join(target_folder, "JPEGImages")):
        os.makedirs(os.path.join(target_folder, "JPEGImages"))
        os.makedirs(os.path.join(target_folder, "Annotations"))
        os.makedirs(os.path.join(target_folder, "ImageSets", "Main"))

    image_counter = 0
    all_filenames = []
    class_files = {}

    for source_folder in source_folders:
        jpeg_images_path = os.path.join(source_folder, "JPEGImages")
        annotations_path = os.path.join(source_folder, "Annotations")

        for filename in os.listdir(jpeg_images_path):
            src_image = os.path.join(jpeg_images_path, filename)
            barename = filename.split(".")[0]
            dst_image = os.path.join(
                target_folder, "JPEGImages", f"{barename}_{image_counter}.jpg"
            )
            annotation_filename = filename.replace(".jpg", ".xml")
            src_annotation = os.path.join(annotations_path, annotation_filename)

            tree = ET.parse(src_annotation)
            root = tree.getroot()

            objects = root.findall("object")
            for obj in objects:
                name = obj.find("name").text
                if name in excluded_classes:
                    root.remove(obj)

            remaining_objects = root.findall("object")
            if not remaining_objects:
                continue

            dst_image = os.path.join(
                target_folder, "JPEGImages", f"{barename}_{image_counter}.jpg"
            )
            shutil.copy(src_image, dst_image)

            dst_annotation = os.path.join(
                target_folder, "Annotations", f"{barename}_{image_counter}.xml"
            )
            tree.write(dst_annotation)

            new_filename = f"{barename}_{image_counter}"
            all_filenames.append(new_filename)

            for obj in remaining_objects:
                name = obj.find("name").text
                if name not in class_files:
                    class_files[name] = []
                class_files[name].append(new_filename)

            image_counter += 1

    # Calculate the total number of samples after merging
    total_samples = sum(total_class_counts.values())

    # Resample classes based on the provided configuration and proportions
    for cls, config in resampling_config.items():
        method = config["method"]
        proportion = config.get("proportion", None)

        if cls in class_files:
            current_count = len(class_files[cls])
            if proportion is not None:
                target_count = int(proportion * total_samples)
            else:
                if method == "undersample":
                    target_count = min(total_class_counts.values())
                elif method == "oversample":
                    target_count = max(total_class_counts.values())

            if method == "undersample" and current_count > target_count:
                sampled_indices = random.sample(range(current_count), target_count)
                class_files[cls] = [class_files[cls][i] for i in sampled_indices]
            elif method == "oversample" and current_count < target_count:
                oversampled_files = random.choices(
                    class_files[cls], k=target_count - current_count
                )
                class_files[cls].extend(oversampled_files)

    balanced_filenames = set()
    for cls, files in class_files.items():
        balanced_filenames.update(files)

    return list(balanced_filenames)


def split_dataset(filenames, train_ratio=0.9):
    random.shuffle(filenames)
    num_train = int(len(filenames) * train_ratio)
    train_set = filenames[:num_train]
    val_set = filenames[num_train:]
    return train_set, val_set


def write_to_file(file_path, data_list):
    with open(file_path, "w") as file:
        for item in data_list:
            file.write(item + "\n")


if __name__ == "__main__":
    seg_path = r"F:\GitProject\yolov7-pytorch\VOCdevkit\seg"
    dirlist = os.listdir(seg_path)
    source_folders = []  # 添加你的源文件夹路径
    # for foldername in dirlist:
    #     inputdir = os.path.join(seg_path, foldername)
    #     source_folders.append(inputdir)
    source_folders.append(r"F:\GitProject\yolov7-pytorch\VOCdevkit\seg\uv")
    source_folders.append(r"F:\GitProject\yolov7-pytorch\VOCdevkit\furtrainimg")
    voc2025dir = r"F:\GitProject\yolov7-pytorch\VOCdevkit\251003seg"
    for foldername in os.listdir(voc2025dir):
        source_folders.append(os.path.join(voc2025dir, foldername))

    target_folder = r"F:\GitProject\yolov7-pytorch\VOCdevkit\merge"  # 目标文件夹路径

    excluded_classes = EXCLUDED_CLASSES

    # 统计所有源文件夹中的类别数量
    total_class_counts = {}
    for source_folder in source_folders:
        annotations_path = os.path.join(source_folder, "Annotations")
        class_counts = count_classes_in_annotations(
            annotations_path, excluded_classes=excluded_classes
        )
        for cls, count in class_counts.items():
            if cls in total_class_counts:
                total_class_counts[cls] += count
            else:
                total_class_counts[cls] = count

    print("Total class counts:", total_class_counts)

    # 定义要进行欠采样和过采样的类别及比例
    resampling_config = {
        # "1": {
        #     "method": "undersample",
        #     "proportion": 0.5,
        # },  # 将类别 '1' 欠采样到总样本数的 10%
        # "2": {
        #     "method": "oversample",
        #     "proportion": 0.2,
        # },  # 将类别 '2' 过采样到总样本数的 20%
        # "3": {
        #     "method": "undersample",
        #     "proportion": 0.1,
        # },  # 将类别 '1' 欠采样到总样本数的 10%
        # "white": {
        #     "method": "undersample",
        #     "proportion": 0.2,
        # },  # 将类别 '1' 欠采样到总样本数的 10%
    }

    all_filenames = merge_voc_folders(
        source_folders,
        target_folder,
        resampling_config,
        total_class_counts,
        excluded_classes=excluded_classes,
    )

    train_set, val_set = split_dataset(all_filenames)

    train_txt_path = os.path.join(target_folder, "ImageSets", "Main", "train.txt")
    val_txt_path = os.path.join(target_folder, "ImageSets", "Main", "val.txt")

    write_to_file(train_txt_path, train_set)
    write_to_file(val_txt_path, val_set)
